import os
import pandas as pd
from sqlalchemy import create_engine


class DataUtil:
    def __init__(self):
        pass

    def get_db_connect(self):
        try:
            engine = create_engine('mysql+pymysql://root:123456@localhost:3306/test?charset=utf8')
            print('数据库连接成功：get_db_connect()')
            return engine
        except Exception as e:
            print('数据库连接失败：get_db_connect()——{}'.format(e))
            return None

    # 读取数据从数据库
    def query_data_from_db(self, sql):
        try:
            engine = self.get_db_connect()
            if engine is None:
                print('数据库连接失败：query_data_from_db()')
                return None
            data = pd.read_sql(sql, engine)
            return data
        except Exception as e:
            print('数据查询失败：query_data_from_db()——{}'.format(e))
            return None

    # 保存数据到数据库
    def save_data_to_db(self, data, table_name, if_exists='replace', index=False):  # index属性，如果为true,则会在最左侧出现一列索引
        try:
            engine = self.get_db_connect()
            if engine is None:
                print('数据库连接失败：save_data_to_db()')
                return None
            # data.columns = ['ip', 'province', 'city', 'access', 'mobile', 'browser', 'gender', 'age', 'method', 'url',
            #                 'code']  写入数据库默认列名为0,1,2,3,4,5,6,7,8,9,10,11，所以可以使用该方法指定列名
            data.to_sql(table_name, engine, if_exists=if_exists, index=index)
            print('数据保存成功：save_data_to_db()')
        except Exception as e:
            print('数据保存失败：save_data_to_db()——{}'.format(e))

    # 读取csv文件
    def read_data_from_csv(self, file_path):
        try:
            if not os.path.exists(file_path):
                print('文件不存在：{}'.format(file_path))
                return None
            # header=None 防止默认将第一行数据当作表头,通过DataFrame.columns可以为指定数据添加列名
            data = pd.read_csv(filepath_or_buffer=file_path, sep=',', encoding='utf-8', header=None)
            print('数据读取成功：read_data_from_csv()')
            return data
        except Exception as e:
            print('读取csv文件失败：read_data_from_csv()——{}'.format(e))
            return None

    # 保存数据到csv文件
    def save_data_to_csv(self, data, file_path, index=False, header=False):
        try:
            if os.path.exists(file_path):
                print('文件内容已被新数据覆盖：{}'.format(file_path))
            # 如果文件已经存在，会默认重写文件,也可以选择对文件进行名称更改后备份，或者增加人机交互
            data.to_csv(file_path, index=index, header=header)
            print('数据写入{}文件成功：save_data_to_csv()'.format(file_path))
        except Exception as e:
            print('数据保存失败：save_data_to_csv()——{}'.format(e))

    # 对数据进行清洗和缺失值处理
    def clean_data(self, data):
        try:
            # 为数据添加列名
            data.columns = ['ip', 'province', 'city', 'access_time', 'mobile', 'browser', 'gender', 'age', 'method',
                            'url', 'status_code']

            # 1. 处理缺失值
            # 对于地理位置信息，使用"未知"填充
            fill_values = {
                'province': '未知',
                'city': '未知',
                'mobile': '未知',
                'browser': '未知',
                'gender': '未知',
                'age': data['age'].mean()
            }
            data.fillna(value=fill_values, inplace=True)

            # 2. 处理重复数据
            data.drop_duplicates(inplace=True)

            # 3. 数据类型转换
            # 将age转换为整数类型
            data['age'] = data['age'].astype(int)

            # 将status_code转换为整数类型
            data['status_code'] = data['status_code'].astype(int)

            # 4. 处理异常值
            # 年龄应该在合理范围内(0-120)
            data = data[(data['age'] >= 0) & (data['age'] <= 120)]

            # 5. 处理时间格式
            # 确保access_time是datetime类型
            data['access_time'] = pd.to_datetime(data['access_time'])

            print('数据清洗完成：clean_data()')
            return data
        except Exception as e:
            print('数据清洗失败：clean_data()——{}'.format(e))
            return None
